"""
训练模型
"""
import os
import torch
import torch.nn.functional as F
from torch.optim import Adam
from tqdm import tqdm

import config
from dnn.sort.siamese import SiameseNetWork
from dnn.sort.dataset import data_loader

# 创建模型和优化器
model = SiameseNetWork().to(config.device)
optimizer = Adam(model.parameters(), lr=1e-3)
# 模型的加载
if os.path.exists(config.sort_save_model_path):
    model.load_state_dict(torch.load(config.sort_save_model_path))
    optimizer.load_state_dict(torch.load(config.sort_save_optimizer_path))


def train(epoch):
    bar = tqdm(enumerate(data_loader), total=len(data_loader), ascii=True, desc="dnnsort模型训练")
    for idx, (input1, input2, target) in bar:
        # inputX: [batch_size, max_len]
        input1 = input1.to(config.device)
        input2 = input2.to(config.device)
        # target: [batch_size,]
        target = target.to(config.device)
        # print(target)

        # 1. 预测结果
        # [batch_size, 2]
        pre = model(input1, input2)

        # 2. 计算损失
        loss = F.nll_loss(pre, target)

        # 3. 梯度置零
        optimizer.zero_grad()

        # 4. 反向传播
        loss.backward()

        # 5. 更新参数
        optimizer.step()

        bar.set_description("epoch:{} idx:{} loss:{:.6f}".format(epoch, idx, loss.item()))

        if idx % 100 == 0:
            torch.save(model.state_dict(), config.sort_save_model_path)
            torch.save(optimizer.state_dict(), config.sort_save_optimizer_path)
